package com.ltx.aishare.config;

import com.ltx.aishare.service.ToolsService;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.service.UserMessage;
import jakarta.annotation.Resource;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * @author 罗添煦
 * @description LangChain4j配置类
 * @create 2025-04-20 下午5:35
 */
@Configuration
public class AiConfig {

    @Resource
    private ToolsService toolsService;

    /**
     * 创建语言模型实例
     *
     * @param chatLanguageModel          非流式响应模型
     * @param streamingChatLanguageModel 流式响应模型
     * @return Assistant
     */
    @Bean
    public Assistant assistant(ChatLanguageModel chatLanguageModel, StreamingChatLanguageModel streamingChatLanguageModel) {
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .chatMemory(chatMemory)
                .build();
    }

    /**
     * 创建语言模型实例工具类function-call
     *
     * @param chatLanguageModel          非流式响应模型
     * @param streamingChatLanguageModel 流式响应模型
     * @return Assistant
     */
    @Bean
    public AssistantTool assistantTool(ChatLanguageModel chatLanguageModel, StreamingChatLanguageModel streamingChatLanguageModel,
                               ToolsService toolsService) {
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

        return AiServices.builder(AssistantTool.class)
                .chatLanguageModel(chatLanguageModel)
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .tools(toolsService)
                .chatMemory(chatMemory)
                .build();
    }

    @Bean
    public AssistantUnique assistantUnique(ChatLanguageModel chatLanguageModel, StreamingChatLanguageModel streamingChatLanguageModel) {
        return AiServices.builder(AssistantUnique.class)
                .chatLanguageModel(chatLanguageModel)
                .streamingChatLanguageModel(streamingChatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder()
                        .maxMessages(10)
                        .id(memoryId).build())
                .build();
    }

    /**
     * 语言模型接口
     */
    public interface Assistant {
        String chat(String prompt);

        // 流式响应输出
        TokenStream stream(String prompt);
    }

    /**
     * 语言模型接口唯一Id进行对话隔离
     */
    public interface AssistantUnique {
        String chat(@MemoryId int memoryId, @UserMessage String prompt);

        // 流式响应输出
        TokenStream stream(@MemoryId int memoryId, @UserMessage String prompt);
    }

    /**
     * 语言模型接口
     */
    public interface AssistantTool {
        String chat(String prompt);

        // 流式响应输出
        TokenStream stream(String prompt);
    }
}
